import asyncio
import time
import random
import aiohttp
from typing import List, Dict, Tuple
from collections import defaultdict


class ConcurrentModelTester:
    def __init__(self, api_url: str, api_key: str = None, max_workers: int = 10):
        """
        初始化并发测试器
        
        :param api_url: 大模型API的URL
        :param api_key: API密钥(如果需要)
        :param max_workers: 最大并发数
        """
        self.api_url = api_url
        self.api_key = api_key
        self.max_workers = max_workers
        self.session = None
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}" if api_key else ""
        }
    
    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self
    
    async def __aexit__(self, exc_type, exc, tb):
        await self.session.close()
    
    async def call_model(self, prompt: str, session: aiohttp.ClientSession) -> Tuple[float, Dict]:
        """
        调用大模型API
        
        :param prompt: 输入的提示词
        :param session: aiohttp会话
        :return: (响应时间, 响应数据)
        """
        payload = {
            "model": "qwen-omni-2.5",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.7
        }
        
        start_time = time.time()
        try:
            async with session.post(
                self.api_url,
                json=payload,
                headers=self.headers
            ) as response:
                response_data = await response.json()
                elapsed = time.time() - start_time
                return elapsed, response_data
        except Exception as e:
            elapsed = time.time() - start_time
            return elapsed, {"error": str(e)}
    
    async def run_concurrent_test(self, prompts: List[str], concurrency: int) -> Dict:
        """
        运行并发测试
        
        :param prompts: 提示词列表
        :param concurrency: 并发数
        :return: 测试结果统计
        """
        if not self.session:
            raise RuntimeError("Session not initialized. Use async with.")
        
        semaphore = asyncio.Semaphore(concurrency)
        results = []
        
        async def worker(prompt: str):
            async with semaphore:
                return await self.call_model(prompt, self.session)
        
        start_time = time.time()
        tasks = [worker(prompt) for prompt in prompts]
        results = await asyncio.gather(*tasks)
        total_time = time.time() - start_time
        
        # 统计结果
        stats = {
            "total_requests": len(prompts),
            "concurrency": concurrency,
            "total_time": total_time,
            "requests_per_second": len(prompts) / total_time,
            "response_times": [],
            "success_count": 0,
            "error_count": 0,
            "errors": defaultdict(int)
        }
        
        for elapsed, response in results:
            stats["response_times"].append(elapsed)
            if "error" in response:
                stats["error_count"] += 1
                stats["errors"][response["error"]] += 1
            else:
                stats["success_count"] += 1
        
        # 计算响应时间统计
        if stats["response_times"]:
            stats["avg_response_time"] = sum(stats["response_times"]) / len(stats["response_times"])
            stats["min_response_time"] = min(stats["response_times"])
            stats["max_response_time"] = max(stats["response_times"])
            stats["p95_response_time"] = sorted(stats["response_times"])[int(len(stats["response_times"]) * 0.95)]
        else:
            stats["avg_response_time"] = 0
            stats["min_response_time"] = 0
            stats["max_response_time"] = 0
            stats["p95_response_time"] = 0
        
        return stats


async def main():
    # 配置参数
    API_URL = "https://api.example.com/v1/chat/completions"  # 替换为实际API地址
    API_KEY = "your-api-key"  # 替换为你的API密钥
    TEST_PROMPTS = [
        "解释量子计算的基本原理",
        "写一首关于春天的诗",
        "如何学习Python编程",
        "世界十大著名建筑有哪些",
        "用Python实现快速排序算法",
        "人工智能的未来发展趋势",
        "如何提高英语口语能力",
        "解释区块链技术的工作原理",
        "推荐几本好的科幻小说",
        "健康饮食的建议有哪些"
    ] * 10  # 重复10次，共100个请求
    
    CONCURRENCY_LEVELS = [1, 5, 10, 20, 30, 50]  # 测试不同的并发级别
    
    # 运行测试
    async with ConcurrentModelTester(API_URL, API_KEY, max_workers=max(CONCURRENCY_LEVELS)) as tester:
        print("开始并发测试...")
        print(f"总请求数: {len(TEST_PROMPTS)}")
        print("=" * 80)
        
        for concurrency in CONCURRENCY_LEVELS:
            print(f"\n测试并发数: {concurrency}")
            stats = await tester.run_concurrent_test(TEST_PROMPTS, concurrency)
            
            print(f"总时间: {stats['total_time']:.2f}秒")
            print(f"每秒请求数: {stats['requests_per_second']:.2f}")
            print(f"平均响应时间: {stats['avg_response_time']:.2f}秒")
            print(f"最小响应时间: {stats['min_response_time']:.2f}秒")
            print(f"最大响应时间: {stats['max_response_time']:.2f}秒")
            print(f"P95响应时间: {stats['p95_response_time']:.2f}秒")
            print(f"成功请求: {stats['success_count']}")
            print(f"失败请求: {stats['error_count']}")
            
            if stats['error_count'] > 0:
                print("\n错误统计:")
                for error, count in stats['errors'].items():
                    print(f"  {error}: {count}次")
            
            print("-" * 60)


if __name__ == "__main__":
    asyncio.run(main())
